# Set pip source
RUN mkdir /root/.pip ;\
    echo "[global]" > /root/.pip/pip.conf ;\
    echo "index-url=https://mirrors.huaweicloud.com/repository/pypi/simple" >> /root/.pip/pip.conf ;\
    echo "trusted-host=repo.huaweicloud.com" >> /root/.pip/pip.conf ;\
    echo "timeout=120" >> /root/.pip/pip.conf

path_resource=/data
version_torch=v2.1
version_torch=v1.11.0
py_version=python3.7
compiler=clang # clang or gcc

#* 0表示不使用，
#* 1 表示生成 lto 优化，
#* 2 表示使用 pgo 第一次编译
#* 3 表示使用 pgo 第二次编译
export flag_use_pgo=1

func_check_pytorch(){
    cd $path_resource/pytorch
    if [[ ! -d ./pytorch-$version_torch ]];then
        echo "尝试寻找包"
        if [[ ! -f pytorch-$version_torch.tar.gz ]];then
            echo -e "\033[33m pytorch-$version_torch.tar.gz 不存在，尝试下载 \033[0m"
            wget http://mirror.ghproxy.com/https://github.com/pytorch/pytorch/releases/download/$version_torch/pytorch-$version_torch.tar.gz
        fi
        tar -zxvf pytorch-$version_torch.tar.gz
    fi

    if [[ ! -d /tmp/pytorch ]];then
        cd $path_resource/pytorch
        echo "复制pytorch${version_torch}源码到 /tmp/pytorch"
        mkdir -p /tmp/pytorch
        cp -r ./pytorch-$version_torch/* /tmp/pytorch
    fi
    cd /tmp/pytorch
    cat /tmp/pytorch/CMakeLists.txt | grep "string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2" # 检查

    if [[ $compiler == clang ]];then
        flag_flto=thin
    else
        flag_flto=auto
    fi
    if [[ $flag_use_pgo == 1 ]];then
        string_new_cmd="-mcpu=tsv110\ -fPIC\ -flto=$flag_flto\ -fuse-ld=lld"
    fi
    if [[ $flag_use_pgo == 2 ]];then
        string_new_cmd="-mcpu=tsv110\ -fPIC\ -flto=$flag_flto\ -fuse-ld=lld -fprofile-generate=/root/pgo_profile"
    fi
    if [[ $flag_use_pgo == 3 ]];then
        string_new_cmd="-mcpu=tsv110\ -fPIC\ -flto=$flag_flto\ -fuse-ld=lld -fprofile-use=/root/pgo_profile/pgo.profdata"
    fi

    if [[ x${flag_use_pgo} != x0 ]];then
        sed -i "s@string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2\ .*@string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2\ $string_new_cmd\")@g" /tmp/pytorch/CMakeLists.txt
        cat /tmp/pytorch/CMakeLists.txt | grep "string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2" # 检查
        #! 注释掉一个编译选项
        string_new_cmd1="append_cxx_flag_if_supported(\"-Werror=cast-function-type\" CMAKE_CXX_FLAGS)"
        string_new_cmd2="string(APPEND CMAKE_CXX_FLAGS \" -Werror=cast-function-type\")"
        sed -i "s@^.*${string_new_cmd1}.*\$@\ \ #\ ${string_new_cmd1}@g" /tmp/pytorch/CMakeLists.txt #! 新版本的 torch 代码
        sed -i "s@^.*${string_new_cmd2}.*\$@\ \ #\ ${string_new_cmd2}@g" /tmp/pytorch/CMakeLists.txt #! 旧版本的 torch 代码
        cat /tmp/pytorch/CMakeLists.txt | grep -n 'Werror=cast-function-type'
    fi
    /usr/local/bin/$py_version -m pip install pyyaml
    #! torch1版本下使用 clang 的话，需要关闭掉一个第三方库

    if [[ x${version_torch:1:1} == x1 ]];then
        export BUILD_BREAKPAD=0
        export USE_BREAKPAD=0
    fi
    /usr/local/bin/$py_version -m pip install -r requirements.txt
    /usr/local/bin/$py_version setup.py bdist_wheel 2>&1 | tee /data/log/pytorch_build.log
    echo "编译后的位置在 dist/torch*.whl"
    # 编译 torch 后安装依赖
    cd /tmp/pytorch
    export USE_XNNPACK=0
    /usr/local/bin/$py_version -m pip install dist/torch*.whl
}


string_new_cmd="-mcpu=tsv110\ -fGNU-compatibility\ -flto=thin\ -fuse-ld=lld\ "
sed -i "s@^SET(CMAKE_CXX_FLAGS_RELEASE\ \"-O2.*@SET(CMAKE_CXX_FLAGS_RELEASE\ \"-O2\ ${string_new_cmd}\")@g" /tmp/pta/CMakeLists.txt
cat /tmp/pta/CMakeLists.txt | grep -n "SET(CMAKE_CXX_FLAGS_RELEASE\ \"-O2"  # 检查
cd /tmp/pta
ln -sf /usr/local/bin/$py_version /usr/bin/python3
alias python3=/usr/local/bin/python3.7
bash ci/build.sh --python=3.7



